import asyncio
import logging
import websockets
import json
from asr_server import ASRServer, ASRConfig

# ASR配置 - 替换为你的实际配置
CONFIG  = ASRConfig(
        app_id="f4b3bfd8",
        api_key="3f6b034a6397fab4d4478345f42340e2",
        api_secret="MDdhM2I3NjM0NjViOTVlY2JkZjQ5NTZj"
    )

class ASRWSServer:
    def __init__(self,config:ASRConfig,client_timeout=3):
        # self.asr_server = ASRServer.create_asr_server(CONFIG)
        self.clients = set()
        self.client_timeout = client_timeout
        self.config = config

    async def handle_client(self, websocket, path):
        """处理WebSocket连接和音频数据"""
        self.clients.add(websocket)
        client_id = id(websocket)
        server = ASRServer.create_asr_server(CONFIG,client_id)
        print(f"New client connected. Total clients: {len(self.clients)}")
        
        async def callback(result):
            response = {
                "text": result.bestMatchText,
                "full_text": result.full_text,
                "status": result.status
            }
            await websocket.send(json.dumps(response))

        # 实时处理识别结果
        if not server.ws:
            server.create_asr_connection()
        # 创建任务处理识别结果
        asyncio.create_task(self._process_results(server, callback))
        try:
            while True:
                message = await asyncio.wait_for(
                    websocket.recv(),
                    timeout=self.client_timeout  # 500ms timeout
                )
                if isinstance(message, bytes) and len(message)>0:
                    await server.write_audio(message)
        except asyncio.TimeoutError:
            print("ws timeout, will close connection")
        except Exception as e:
            print(f"Client error: {e}:{str(e)}")
        finally:
            self.clients.remove(websocket)
            print(f"Client disconnected. Remaining clients: {len(self.clients)}")

    async def _process_results(self, server, callback):
        """处理ASR识别结果"""
        async for result in server.process_results(callback):
            pass

    async def run(self, host="0.0.0.0",port=8765,ssl=False,certfile=None,keyfile=None):
        """启动WebSocket服务器"""
        if ssl:
            ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
            ssl_context.load_cert_chain(certfile, keyfile)
            async with websockets.serve(self.handle_ws_client, host, port, ssl=ssl_context):
                await asyncio.Future()
            logging.info(f"Starting test server on wss://{host}:{port}")
        else:
            logging.info(f"Starting test server on ws://{host}:{port}")
            async with websockets.serve(self.handle_client, host, port):
                await asyncio.Future()  # 永久运行  

if __name__ == "__main__":
    server = ASRWSServer(CONFIG)
    asyncio.run(server.run())
